from neo4j import GraphDatabase
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
import numpy as np
from neo4j.data import Record
from nodevec.neo4j_graph_sql import *
from nodevec.fastRP_sql import *
# Global defaults, some based on our demo and some on the algo defaults.
DEFAULT_URI = "bolt://localhost:7687"
DEFAULT_USER = "neo4j"
DEFAULT_PASS = "123456"
DEFAULT_PROP = "clusterId"
DEFAULT_K = 10



def get_company_ids(driver):
    cypter = '''
    call {Match (n:`企业`) return n} return id(n)'''
    ids = []
    with driver.session() as session:
        results = session.run(cypter)
        for result in results:
            ids.append(result[0])
    return ids

def extract_embeddings(driver,company_ids):
    embeddings = []
    with driver.session() as session:
        results = session.run(NODE2VEC_fastPR_vec)
        for result in results:
            if result["nodeId"] not in company_ids:
                continue
            embeddings.append(result)
    return embeddings




def _update_tx(tx, cypher, **kwargs):
    result = tx.run(cypher, kwargs)
    return result.consume()

def update_clusters(driver, clusterResults):
    """
    Given a list of dicts with "nodeId" string and a "valueMap" dict, update
    the graph by setting the properties from the "valueMap" on each node.
    """
    print("Updating graph...")
    with driver.session() as session:
        result = session.write_transaction(_update_tx, UPDATE_CYPHER, updates=clusterResults)
        print("...update complete: {}".format(result.counters))

import os
driver = GraphDatabase.driver(DEFAULT_URI, auth=(DEFAULT_USER, DEFAULT_PASS))

company_ids = get_company_ids(driver)
embeddings = extract_embeddings(driver,company_ids)




clusters = kmeans_cluster(embeddings, k=-1, clusterProp=DEFAULT_PROP)
update_clusters(driver, clusters)
driver.close()




